import os
import pickle
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error

# 设置设备（有GPU用GPU，否则CPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 确保负号正常显示（matplotlib可选）
import matplotlib.pyplot as plt
plt.rcParams['axes.unicode_minus'] = False

print("开始加载数据并准备预测模型...")

# 读取数据
data = pd.read_csv('historical_electricity_data.csv')
data['timestamp'] = pd.to_datetime(data['timestamp'])
print("数据基本信息：")
data.info()

prediction_vars = ['load_kW', 'grid_price', 'wind_price', 'solar_price']
var_names = {
    'load_kW': '电力负荷(kW)',
    'grid_price': '国网电价(元/kWh)',
    'wind_price': '风电电价(元/kWh)',
    'solar_price': '太阳能电价(元/kWh)'
}

look_back = 24  # 时间步长

# 定义LSTM模型
class LSTMModel(nn.Module):
    def __init__(self, input_size=1, hidden_size=50, num_layers=2):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 25),
            nn.Linear(25, 1)
        )

    def forward(self, x):
        out, _ = self.lstm(x)
        out = out[:, -1, :]  # 取最后时间步输出
        out = self.fc(out)
        return out

# 构造数据集
def create_dataset(dataset, look_back=24):
    X, Y = [], []
    for i in range(len(dataset) - look_back):
        X.append(dataset[i:i+look_back])
        Y.append(dataset[i+look_back])
    return np.array(X), np.array(Y)

# 训练参数
num_epochs = 50
batch_size = 32
learning_rate = 0.001

models = {}
scalers = {}
rmse = {}

for var in prediction_vars:
    print(f"\n训练变量：{var_names[var]}")

    var_data = data[var].values.reshape(-1,1)
    scaler = MinMaxScaler(feature_range=(0,1))
    scaled_data = scaler.fit_transform(var_data)
    scalers[var] = scaler

    X, y = create_dataset(scaled_data, look_back)

    # 80%训练集
    train_size = int(len(X)*0.8)
    X_train, X_test = X[:train_size], X[train_size:]
    y_train, y_test = y[:train_size], y[train_size:]

    # 转 torch tensor 并放到设备上
    X_train = torch.tensor(X_train, dtype=torch.float32).to(device)
    y_train = torch.tensor(y_train, dtype=torch.float32).to(device)
    X_test = torch.tensor(X_test, dtype=torch.float32).to(device)
    y_test = torch.tensor(y_test, dtype=torch.float32).to(device)

    model = LSTMModel().to(device)
    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    # 训练循环
    model.train()
    for epoch in range(num_epochs):
        permutation = torch.randperm(X_train.size(0))
        epoch_loss = 0
        for i in range(0, X_train.size(0), batch_size):
            indices = permutation[i:i+batch_size]
            batch_x, batch_y = X_train[indices], y_train[indices]

            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs.squeeze(), batch_y.squeeze())
            loss.backward()
            optimizer.step()

            epoch_loss += loss.item()

        avg_loss = epoch_loss / (X_train.size(0) / batch_size)
        if (epoch+1) % 10 == 0 or epoch == 0:
            print(f"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.6f}")

    models[var] = model

    # 评估
    model.eval()
    with torch.no_grad():
        preds = model(X_test).squeeze().cpu().numpy()
    preds = scaler.inverse_transform(preds.reshape(-1,1))
    actual = scaler.inverse_transform(y_test.cpu().numpy().reshape(-1,1))
    rmse_val = np.sqrt(mean_squared_error(actual, preds))
    rmse[var] = rmse_val
    print(f"{var_names[var]} 测试集RMSE: {rmse_val:.4f}")

# 保存模型和scaler
model_dir = "saved_models"
os.makedirs(model_dir, exist_ok=True)

for var in prediction_vars:
    model_path = os.path.join(model_dir, f"{var}_lstm_model.pt")
    torch.save(models[var].state_dict(), model_path)
    print(f"已保存模型: {model_path}")

    scaler_path = os.path.join(model_dir, f"{var}_scaler.pkl")
    with open(scaler_path, 'wb') as f:
        pickle.dump(scalers[var], f)
    print(f"已保存归一化器: {scaler_path}")
